"""
    给GRR结果表格添加基因的gog_category和gene_description
"""

import os
import argparse
import pandas as pd 


def add_gene_info(cluster_table, gene_info_file, output_file):
    # 读入基因簇结果表格
    cluster_df = pd.read_excel(cluster_table)
    cluster_title = cluster_df.columns.tolist()
    # print(cluster_title.index("gene_id"))
    # print(cluster_title)
    
    # 读入基因对应描述信息表格
    if os.path.exists(gene_info_file):
        gene_info = pd.read_csv(gene_info_file, sep="\t", dtype="str")
        gene_info_title = gene_info.columns.tolist() 
        # print(gene_info_title) 
        # print(gene_info)

        # 基因id作为键，对应行作为值
        gene_info_dict = {}
        for row in gene_info.values.tolist():
            gene = row[0]
            gene_info_dict[gene] = row
            
        result = []
        for row in cluster_df.values.tolist():
            tmp_row = row 
            gene_name = row[7]
            info = ""
            if gene_name in gene_info_dict:
                info = gene_info_dict[gene_name]
            else:
                print(gene_name)

            gene_description = info[3]
            if gene_description == "-":
                gene_description = ""
            go_term = info[4]
            gog_category = info[-1]
            if gog_category == "":
                gog_category = ""
            if str(gog_category).lower().strip() == "function unknown":
                gog_category = ""
                
            # print(gene_description, gog_category, "\n")
            tmp_row = tmp_row + [gene_description, gog_category]
            result.append(tmp_row)
            
        cluster_title = cluster_title + ["gene_description", "gog_category"] 

        result_df = pd.DataFrame(result, columns=cluster_title)
        result_df.to_excel(output_file, index=False)
        
        description_num = len(list(set(result_df["gene_description"].tolist())))
        category_num = len(list(set(result_df["gog_category"].tolist())))
        
        with open("02.description_category_num.txt", mode="a") as file:
            file.write(
                "\n"+output_file + "\t" + 
                str(description_num) + 
                "\t" + str(category_num) + "\n"
            )
    else:
        cluster_df["gene_description"] = cluster_df.shape[0] * [""]
        cluster_df["gog_category"] = cluster_df.shape[0] * [""]
        result_df.to_excel(output_file, index=False)
        
        with open("02.description_category_num.txt", mode="a") as file:
            file.write(
                "\n"+output_file + "\t" + 
                str(0) + 
                "\t" + str(0) + "\n"
            )


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-cluster","--cluster_table")
    parser.add_argument("-gene","--gene_info_file")
    parser.add_argument("-output","--output_file")
    args = parser.parse_args()
    
    return args 


if __name__ == "__main__":
    params = parse_args()
    cluster_table = params.cluster_table
    gene_info_file = params.gene_info_file
    output_file = params.output_file
    try:
        add_gene_info(cluster_table, gene_info_file, output_file)
    except Exception as e:
        print(gene_info_file)